import time
import datetime
import pandas as pd
import tushare as ts
from tools.utils import load_config
from tools.utils import loc_collection
from tools.utils import loc_db


def get_list_future_code(cfg: dict):
    """
    :param cfg: 项目总配置文件，来源于config.yaml,mongodb中database:project_config也有存储
    :return:以tushare中fut_code作为key的字典，value是dict,其中索引ts_code可得到以list输出的当前存续合约
    """
    # 获取每个品种当前上市的合约数，可以从本地mongodb调取，也可以从tushare获取
    trade_contract = loc_collection('project_config', 'trade_contract', cfg)
    if not cfg['update_trade_contract']:
        # print("loading trading contract in database .....")
        output_dic = list(trade_contract.find())[-1]
    else:
        print("loading trading contract in tushare .....")
        output_dic = {}
        pro = ts.pro_api(cfg['tushare'])
        for exchange, e_dic in cfg['exchange'].items():
            exchange_ts = cfg['exchange_map'][exchange]
            today_ = datetime.date.today()
            today = str(today_).replace("-", "")
            exchange_df = pro.fut_basic(exchange=exchange_ts, fut_type=1)
            for chinese_name, v_li in e_dic.items():
                code = v_li[1].split('_')[0].upper()
                print(chinese_name, code)
                fut_df = exchange_df[exchange_df['fut_code'] == code]
                fit_df = fut_df[fut_df['delist_date'] > today]
                df = fit_df.sort_values(by=['delist_date'])
                df.index = list(range(len(df)))
                code_dic = {k: v.to_list() for k, v in df.iteritems()}
                output_dic[code] = code_dic
                print(code, " : ", df['ts_code'].to_list())
        print("update to mongodb database: project config collection: trade contract")
        trade_contract.insert_one(output_dic)
    return output_dic


def get_main_future_contract_map(cfg):
    """
    :param cfg: 项目总配置文件，来源于config.yaml,mongodb中database:project_config也有存储
    :return:DataFrame key是rb_shf的主力合约无。名字，index是每日时间，单元格是合约在t时刻的对应合约。
    """
    main_contract_map = loc_collection('project_config', 'main_contract_map', cfg)
    if cfg['update_main_contract_map']:
        print("removing main_contract_map !!!!")
        main_contract_map.remove()
        pro = ts.pro_api(cfg['tushare'])
        li = []
        for exchange, e_dic in cfg['exchange'].items():
            for chinese_name, v_li in e_dic.items():
                print(chinese_name, v_li)
                ts_code = v_li[1].upper().replace('_', '.')
                df = pro.fut_mapping(ts_code=ts_code)
                df.index = df['trade_date']
                df = df.sort_index()
                se = df['mapping_ts_code']
                se.name = v_li[1]
                li.append(se)
        ddf = pd.concat(li, axis=1)
        ddf = ddf.sort_index()
        ddf['t'] = ddf.index
        v_li = [se.to_dict() for t, se in ddf.iterrows()]
        main_contract_map.insert_many(v_li)
        ddf = ddf.drop(['t'], axis=1)
    else:
        print("loading main contract map from mongo")
        v_li = main_contract_map.find()
        ddf = pd.DataFrame(v_li)
        ddf.index = ddf['t']
        ddf = ddf.drop(['_id', 't'], axis=1)
    return ddf


def update_future_to_mongo():
    today = str(datetime.date.today())
    conf = load_config()
    pro = ts.pro_api(conf['tushare'])
    db = loc_db('tushare_future', conf)
    coll_names = db.list_collection_names()
    for coll_name in coll_names:
        coll = loc_collection('tushare_future', coll_name, conf)
        data = list(coll.find({'update': True}))
        for dic_in_col in data:
            end_t = str(pd.to_datetime(str(dic_in_col['end_t'])))[:10]
            ts_code = dic_in_col['ts_code']
            hour = time.localtime().tm_hour
            print('ent_t', end_t, 'today', today)
            if end_t < today:
            # and hour > 20:  # 更新今日数据,并且在8点以后
                daily_data = pro.fut_daily(ts_code=ts_code)
                daily_data = daily_data.sort_values(['trade_date'])
                dic_data = {k: v.to_list() for k, v in daily_data.items()}
                insert_end_t = dic_data['trade_date'][-1]
                if str(pd.to_datetime(str(insert_end_t)))[:10] > end_t:
                    insert_data = {
                        'ts_code': ts_code,
                        't': [str(pd.to_datetime(str(s)))[:10] for s in dic_data['trade_date']],
                        'start_t': dic_data['trade_date'][0],
                        'end_t': insert_end_t,
                        'update': dic_in_col['update'],
                    }
                    del dic_data['ts_code']
                    for k, v in dic_data.items():
                        insert_data[k] = v
                    coll.remove({'ts_code': ts_code})  # 删除老数据
                    coll.insert(insert_data)
                    print(f'更新{ts_code}到{insert_end_t}')
                    time.sleep(0.5)
                else:
                    print(f'最新数据仍是{end_t},不需要更新！！！')
                    time.sleep(0.34)
            else:
                print('不需要更新', ts_code)


if __name__ == "__main__":
    update_future_to_mongo()
    # cfg = load_config()
    # cfg['update_main_contract_map'] = True
    # get_main_future_contract_map(cfg)
    print(777)
